Add unified encoder pytorch implementation #251

CeliaBenquet · 2025-05-01T14:36:20Z

This PR adds a PyTorch implementation of a unified CEBRA encoder, which is composed of:

A new sampling scheme that samples across all sessions so that they can be aligned on the neuron axis to train a single encoder.
A unified Dataset and Loader, adapted to the new sampling scheme.
A unified Solver that considers multiple sessions to be aligned at inference.
A new masked modeling training option, with different types of masking.

🚧 A preprint is pending "Unified CEBRA Encoders for Integrating Neural Recordings via Behavioral Alignment" by Célia Benquet, Hossein Mirzaei, Steffen Schneider, Mackenzie W. Mathis.

…ional models in _transform

…est accordingly

add headers to new files

- header

MMathisLab · 2025-05-26T10:23:42Z

cebra/data/single_session.py

-            positive=self[index.positive],
-            negative=self[index.negative],
-            reference=self[index.reference],
+            positive=self.apply_mask(self[index.positive]),


quick sanity check; this is backwards compatable? @CeliaBenquet

backward compatible yes, I added a check on if the function doesn't exist for people who might want to use the adapt functionality on an older model, good catch.

cebra/datasets/__init__.py

cebra/integrations/sklearn/cebra.py

- header

MMathisLab

thanks @CeliaBenquet ! I went through and left comments for disucssion

MMathisLab · 2025-05-26T10:25:28Z

cebra/models/decoders.py

should we put this into integrations vs. models? Models to me is encoders only cc @stes

We currently have some decoders here, although these are sklearn specific.

I think this module here is fine, at least right now I dont see a better place in the codebase to put them in. An argument to leave them here would be that they are an "extension" of the encoders we train, plus they are "raw" torch objects, which we currently all collected in cebra.models.

I dont have a strong opinion, just don't see where they would fit better... In integrations, we currently have only "standalone" helper functions, which these aren't.

@CeliaBenquet where are these decoders used around the codebase? and how are they trained?

okay, I see

cebra/models/model.py

cebra/solver/multi_session.py

cebra/solver/single_session.py

cebra/models/decoders.py

MMathisLab · 2025-05-26T10:34:59Z

cebra/models/decoders.py

@@ -0,0 +1,38 @@
+import torch.nn as nn


why not have decoders somewhere like integrations? to me model is the encoders only cc @stes

and this :D

stes

Looks good overall; left some comments!

Implementation of the Mixin class for the masking: If I understood correctly, the only change is that this apply_mask function is applied after loading a batch. This seems to be a change that could be minimally invasively applied not in the dataset, but actually in the data loader. Is there a good case why the datasets themselves need to be modified?
Discussion on where to place the decoders: currently in cebra.models.decoders; are the decoders useful as "standalone" models? where are they currently used? based on that we could determine if we move them e.g. as standalone to integrations
see other comments; mostly on class design, removing duplicated code, etc.

cebra/data/masking.py

stes · 2025-05-28T01:35:40Z

cebra/data/multi_session.py

+        if hasattr(self, "apply_mask"):
+            batch = [
+                cebra_data.Batch(
+                    reference=self.apply_mask(
+                        session[index.reference[session_id]]),
+                    positive=self.apply_mask(
+                        session[index.positive[session_id]]),
+                    negative=self.apply_mask(
+                        session[index.negative[session_id]]),
+                    index=index.index,
+                    index_reversed=index.index_reversed,
+                ) for session_id, session in enumerate(self.iter_sessions())
+            ]
+        else:
+            batch = [
+                cebra_data.Batch(
+                    reference=session[index.reference[session_id]],
+                    positive=session[index.positive[session_id]],
+                    negative=session[index.negative[session_id]],
+                    index=index.index,
+                    index_reversed=index.index_reversed,
+                ) for session_id, session in enumerate(self.iter_sessions())
+            ]
+        return batch


can we convert this if/else statement into a subclass

This is a backward compatibility check for old models, I don't know if it's worth it no? ideally we don't have it I added after Mackenzie's comment

Under which circumstance would the apply_mask function be missing?

cebra/data/multi_session.py

stes · 2025-05-28T01:38:21Z

cebra/data/single_session.py

+        if hasattr(self, "apply_mask"):
+            # If the dataset has a mask, apply it to the data.
+            batch = Batch(
+                positive=self.apply_mask(self[index.positive]),
+                negative=self.apply_mask(self[index.negative]),
+                reference=self.apply_mask(self[index.reference]),
+            )
+        else:
+            batch = Batch(
+                positive=self[index.positive],
+                negative=self[index.negative],
+                reference=self[index.reference],
+            )
+        return batch


see above; a better way to implement this is by having the masking simply override the load_batch function, vs. introducing this if/else logic.

This is a backward compatibility check for old models, I don't know if it's worth it no? ideally we don't have it I added after Mackenzie's comment

cebra/datasets/demo.py

docs/source/api/pytorch/helpers.rst

tests/test_data_masking.py

tests/test_solver.py

cebra/data/masking.py

cebra/data/base.py

MMathisLab

lgtm! Just the one comment on kwargs seems critical to decide

gonlairo and others added 30 commits August 23, 2024 13:54

first proposal for batching in tranform method

283de06

first running version of padding with batched inference

202e379

start tests

1f1989d

add pad_before_transform to fit function and add support for convolut…

8665660

…ional models in _transform

remove print statements

8d5b114

first passing test

32c5ecd

add support for hybrid models

9928f63

rewrite transform in sklearn API

be5630a

baseline version of a torch.Datset

1300b20

move batching logic outside solver

bc6af24

move functionality to base file in solver and separate in functions

ec377b9

add test_select_model for single session

6f9ca98

add checks and test for _process_batch

fbe7eb4

add test_select_model for multisession

463b0f8

make self.num_sessions compatible with single session training

5219171

improve test_batched_transform_singlesession

f9bd1a6

make it work with small batches

e23a7ef

make test with multisession work

19c3f87

change to torch padding

87bebac

add argument to sklearn api

f0303e0

add torch padding to _transform

8c8be85

convert to torch if numpy array as inputs

59df402

add distinction between pad with data and pad with zeros and modify t…

1aadc8b

…est accordingly

differentiate between data padding and zero padding

bc8ee25

remove float16

5e7a14c

change argument position

928d882

clean test

07bac1c

clean test

0823b54

Fix warning

9fe3af3

Improve modularity remove duplicate code and todos

b417a23

Adapt unified code to get_model method

32fae46

CeliaBenquet mentioned this pull request May 20, 2025

Add path to unified CEBRA demo icon #252

Open

Merge branch 'AdaptiveMotorControlLab:main' into unified-cebra

e7bedff

CeliaBenquet requested review from stes and MMathisLab May 20, 2025 10:51

MMathisLab mentioned this pull request May 23, 2025

Batched inference CEBRA & padding at the Solver level #168

Merged

MMathisLab changed the base branch from batched-inference-and-padding to main May 23, 2025 13:39

MMathisLab and others added 3 commits May 23, 2025 15:51

Update mask.py

b4caf3a

add headers to new files

Merge branch 'main' into unified-cebra

718f7ca

Update masking.py

bbf0e8f

- header

MMathisLab reviewed May 26, 2025

View reviewed changes

cebra/datasets/__init__.py Outdated Show resolved Hide resolved

MMathisLab reviewed May 26, 2025

View reviewed changes

cebra/integrations/sklearn/cebra.py Outdated Show resolved Hide resolved

Update test_data_masking.py

858b77b

- header

MMathisLab requested changes May 26, 2025

View reviewed changes

CeliaBenquet added 5 commits May 26, 2025 15:10

Implement review comments and fix typos

4d5e9c3

Fix docs errors

4424ba1

Remove np.int typing error

a968768

Fix docstring warning

535cef3

Fix indentation docstrings

8798aa0

stes requested changes May 28, 2025

View reviewed changes

CeliaBenquet added 4 commits May 28, 2025 09:40

Implement review comments

165d641

Fix circular import and abstract method

8e5bd4e

Add maskedmixin to __all__

63d5a7c

Implement extra review comments

d91949f

MMathisLab approved these changes May 28, 2025

View reviewed changes

CeliaBenquet added 3 commits May 28, 2025 20:18

Change masking kwargs as tuple and not dict in sklearn impl

de300a9

Add integrations/decoders.py

fe341e1

Fix typo

276e5a3

MMathisLab requested a review from stes May 28, 2025 23:02

Add unified encoder pytorch implementation #251

Are you sure you want to change the base?

Add unified encoder pytorch implementation #251

Uh oh!

Conversation

CeliaBenquet commented May 1, 2025 • edited by MMathisLab Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

MMathisLab left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

stes left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

MMathisLab left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

CeliaBenquet commented May 1, 2025 •

edited by MMathisLab

Loading